from knext.api.component import TextFileReader, LLMBasedExtractor, KGWriter
from knext.api.auto_prompt import REPrompt
from nn4k.invoker import NNInvoker

model_name = "qwen3"
model_url = "http://supervisor.xinference.svc.cluster.local:9997/v1"
model_key = "your_api_key"
file_input = "input_file.txt"


# 1. 配置大模型服务
llm = NNInvoker.from_config({
    "nn_name": model_name,  # 模型名称
    "openai_api_base": model_url,  # 模型API地址
    "openai_api_key": model_key,  # API密钥
    "openai_max_tokens": 2000,  # 最大生成长度
    "temperature": 0.1  # 温度参数，控制输出随机性
})

# 2. 定义REPrompt抽取配置
re_prompt = REPrompt(
    spg_type_name="Person",  # 目标实体类型（需与Schema中定义一致）
    property_names=["name", "description", "gender", "nickName"],  # 要抽取的属性
    relations=[],  # 要抽取的关系（可选）
    prompt_template="""
    从以下文本中提取人物信息：
    1. 人物名称
    2. 昵称或小名
    3. 性别

    输出格式：JSON，包含name、description、gender、nickName字段。

    文本：{input_text}
    """,
    example=[
        {
            "input": "丁伟伦，是男生，昵称是Allen",
            "output": {
                "name": "丁伟伦",
                "description": "丁伟伦是男生，昵称为Allen",
                "gender": "男",
                "nickName字段": "Allen"
            }
        }
    ],
    constraint_mode=True  # 启用Schema约束，确保输出符合Schema定义
)

# 3. 定义抽取算子
extractor = LLMBasedExtractor(
    llm=llm,
    prompt_ops=[re_prompt],  # 传入REPrompt配置
    batch_size=5,  # 批处理大小
    cache_path="cache.pkl"  # 缓存路径，避免重复调用模型
)

# 4. 定义数据源
source = TextFileReader(
    local_path="/ckpt/file_scanner/" + file_input,  # 替换为实际文件名,路径不用替换
    delimiter="\n\n\n"
)

# 5. 定义输出目标
writer = KGWriter(
    output_path="/ckpt/file_scanner/" +  file_input + ".json",  # 输出文件路径
    format="json"  # 输出格式
)

# 6. 编排组件执行流程
from knext.api.builder import BuilderChain

chain = BuilderChain(components=[source, extractor, writer])
chain.execute()
